package org.hilo.boot.core.db;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;

import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.util.Assert;

/**
 * the routes {@link #getConnection()}
 * calls to one of various target DataSources based on a lookup key. The latter is usually
 * (but not necessarily) determined through some thread-bound transaction context.
 *
 * @author zollty
 * @see #setTargetDataSources
 * @see #setDefaultTargetDataSource
 * @see #determineCurrentLookupKey()
 */
public class DynamicRoutingDataSource extends AbstractDataSource implements InitializingBean {

	private Map<Object, ExtDataSource> targetDataSources;

	private ExtDataSource defaultTargetDataSource;

	private DataSourceLookupStrategy lookupKeyStrategy = new DefaultLookupStrategy();

	/**
	 * Specify the map of target DataSources, with the lookup key as key.
	 * The mapped value can either be a corresponding {@link javax.sql.DataSource}
	 * instance or a data source name String (to be resolved via a
	 * {@link #setDataSourceLookup DataSourceLookup}).
	 * <p>The key can be of arbitrary type; this class implements the
	 * generic lookup process only. The concrete key representation will
	 * be handled by {@link #resolveSpecifiedLookupKey(Object)} and
	 * {@link #determineCurrentLookupKey()}.
	 */
	public void setTargetDataSources(Map<Object, ExtDataSource> targetDataSources) {
		this.targetDataSources = targetDataSources;
	}

	/**
	 * 如果是主从数据源，建议将defaultTargetDataSource设置为master数据源
	 */
	public void setDefaultTargetDataSource(ExtDataSource defaultTargetDataSource) {
		this.defaultTargetDataSource = defaultTargetDataSource;
	}

    public DataSourceLookupStrategy getLookupKeyStrategy() {
        return lookupKeyStrategy;
    }

    public void setLookupKeyStrategy(DataSourceLookupStrategy lookupKeyStrategy) {
        if (lookupKeyStrategy != null) {
            this.lookupKeyStrategy = lookupKeyStrategy;
        }
    }
    
	@Override
	public void afterPropertiesSet() {
		if (this.targetDataSources == null) {
			throw new IllegalArgumentException("Property 'targetDataSources' is required");
		}
		lookupKeyStrategy.setDataSources(targetDataSources);
	}

	@Override
	public Connection getConnection() throws SQLException {
	    ExtDataSource ds = determineTargetDataSource();
	    try {
            return ds.getConnection();
        } catch (SQLException e) {
            lookupKeyStrategy.markConnectionFailure(ds.getName());
            throw e;
        }
	}

	@Override
	public Connection getConnection(String username, String password) throws SQLException {
		ExtDataSource ds = determineTargetDataSource();
        try {
            return ds.getConnection(username, password);
        } catch (SQLException e) {
            lookupKeyStrategy.markConnectionFailure(ds.getName());
            throw e;
        }
	}

	@Override
	@SuppressWarnings("unchecked")
	public <T> T unwrap(Class<T> iface) throws SQLException {
		if (iface.isInstance(this)) {
			return (T) this;
		}
		return determineTargetDataSource().unwrap(iface);
	}

	@Override
	public boolean isWrapperFor(Class<?> iface) throws SQLException {
		return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
	}

	/**
	 * Retrieve the current target DataSource. Determines the
	 * {@link #determineCurrentLookupKey() current lookup key}, performs
	 * a lookup in the {@link #setTargetDataSources targetDataSources} map,
	 * falls back to the specified
	 * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
	 * @see #determineCurrentLookupKey()
	 */
	protected ExtDataSource determineTargetDataSource() {
		Assert.notNull(this.targetDataSources, "DataSource router not initialized");
		Object lookupKey = determineCurrentLookupKey();
		if(lookupKey == null) {
		    return this.defaultTargetDataSource;
		}
		ExtDataSource dataSource = this.targetDataSources.get(lookupKey);
		if (dataSource == null) {
		    return this.defaultTargetDataSource;
		}
		return dataSource;
	}

	/**
	 * Determine the current lookup key. This will typically be
	 * implemented to check a thread-bound transaction context.
	 * <p>Allows for arbitrary keys. The returned key needs
	 * to match the stored lookup key type, as resolved by the
	 * {@link #resolveSpecifiedLookupKey} method.
	 */
    protected Object determineCurrentLookupKey() {
        return lookupKeyStrategy.resolveSpecifiedLookupKey(
                DataSourceKeyHolder.getCurrentKey());
    }

    /**
     * @return the targetDataSources
     */
    public Map<Object,ExtDataSource> getTargetDataSources() {
        return targetDataSources;
    }

}
